package com.fe.class07.shoppingTopk;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;

/**
 * @Description 加强堆结构
 * @Author sosuke :-)
 * @Date 2022/1/9 11:04
 */
public class HeapGreater<T> {
    private final ArrayList<T> heap; //堆存储结构，相当于之前的数组
    private final HashMap<T, Integer> indexMap; //反向索引表，存储元素位置（空间换时间）
    private int heapSize; //堆大小
    private final Comparator<? super T> comparator; //堆中元素的比较器，堆顶是排序最靠前的

    public HeapGreater(Comparator<? super T> comparator) {
        heap = new ArrayList<>();
        indexMap = new HashMap<>();
        heapSize = 0;
        this.comparator = comparator;
    }

    public boolean isEmpty() {
        return heapSize == 0;
    }

    public int size() {
        return heapSize;
    }

    /**
     * 是否存在元素
     *
     * @param obj
     * @return
     */
    public boolean contains(T obj) {
        return indexMap.containsKey(obj);
    }

    /**
     * 查看堆顶元素
     *
     * @return
     */
    public T peek() {
        return heap.get(0);
    }

    /**
     * 向堆中添加元素，维护堆结构，O(logN)
     * @param obj
     */
    public void push(T obj) {
        // 不允许添加重复元素，使用索引表判断
        if (contains(obj)) {
            return;
        }
        // 向堆中添加数据，并维护索引表，向上heapInsert
        heap.add(obj);
        indexMap.put(obj, heapSize);
        heapInsert(heapSize++);
    }

    /**
     * 弹出堆顶元素，堆大小减一，维护堆结构，O(logN)
     * @return
     */
    public T pop() {
        T ans = heap.get(0); // 要返回的堆顶元素
        swap(0, heapSize - 1); // 堆顶元素和最后一个元素交换
        // 在堆中删除最后一个元素，更新索引
        heap.remove(--heapSize);
        indexMap.remove(ans);
        // 从堆顶开始向下heapify
        heapify(0);
        return ans;
    }

    /**
     * 删除堆中任意元素，O(logN)
     * @param obj
     */
    public void remove(T obj) {
        // 元素不在堆中，则返回
        if (!indexMap.containsKey(obj)) return;

        // 记录要删除元素的位置和堆中最后一个元素
        T replaceData = heap.get(heapSize - 1);
        int index = indexMap.get(obj);

        // 在堆中删除最后一个元素,堆大小减一，删除该数据索引
        heap.remove(--heapSize);
        indexMap.remove(obj);

        // 删除元素不是堆中最后一个元素，需要处理
        if (!obj.equals(replaceData)) {
            // 将最后一个元素放到已删除元素的位置，更新索引
            heap.set(index, replaceData);
            indexMap.put(replaceData, index);
            // 重新调整堆（已删除元素位置的元素发生了变化）
            resign(replaceData);
        }
    }

    /**
     * 元素obj发生了变化，要么排序会更靠前(heapInsert)，要么排序会更靠后(heapify)
     * heapInsert 和 heapify只会执行一个
     */
    public void resign(T obj) {
        Integer index = indexMap.get(obj);
        if (null == index) return;
        heapInsert(index);
        heapify(index);
    }

    // 请返回堆上的所有元素
    public List<T> getAllElements() {
        return heap;
    }

    /**
     * 从index开始向上，如果比父节点的排序靠前则交换
     * 移动到0位置，或者干不掉自己的父亲了，停！
     *
     * 比较器返回负数表示第一个数排序靠前
     * @param index
     */
    private void heapInsert(int index) {
        int parentIndex = (index - 1) / 2;
        // index位置的元素排序比父节点元素排序靠前，则交换
        while (comparator.compare(heap.get(index), heap.get(parentIndex)) < 0) {
            swap(index, parentIndex);
            index = parentIndex;
            parentIndex = (index - 1) / 2;
        }
    }

    /**
     * 从index位置开始向下看，如果排序位置在排序靠前的孩子节点之后则交换
     * 停：排序靠前的孩子都不在index位置元素之前；已经没孩子了
     *
     * @param index
     */
    private void heapify(int index) {
        int left = index * 2 + 1;
        while (left < heapSize) { // 左孩子存在
            // 如果右孩子存在且排序位置先于左孩子，则返回右孩子位置，否则返回左孩子位置
            int betterRankIndex = left + 1 < heapSize && comparator.compare(heap.get(left + 1), heap.get(left)) < 0 ? left + 1 : left;
            // 排序靠前的孩子都不在index位置的元素之前，退出
            if (comparator.compare(heap.get(betterRankIndex), heap.get(index)) >= 0) break;

            // index和排序靠前的孩子，要互换
            swap(betterRankIndex, index);
            index = betterRankIndex;
            left = index * 2 + 1;
        }
    }

    private void swap(int i, int j) {
        T o1 = heap.get(i);
        T o2 = heap.get(j);
        // 交换i和j位置的元素，并更新索引表
        heap.set(i, o2);
        heap.set(j, o1);
        indexMap.put(o1, j);
        indexMap.put(o2, i);
    }
}
